#include "Segmentation_GPU.h"
#include <iostream>

Segmentation_GPU::Segmentation_GPU(int width, int height):
bw(height, width),
original(height, width){
	this->width = width;
	this->height = height;
	initMemory();
}


void Segmentation_GPU::initMemory(){
	cudaMalloc(&clusterMap, width * height * sizeof(int));
	cudaMalloc(&clusterMap_C, width * height * sizeof(float4));
	resultMap = new int[width * height];
}


cv::Mat_<int>* Segmentation_GPU::segmentImage(xn::ImageMetaData* imageMD, float THRESHOLD){
	this->THRESHOLD = THRESHOLD;
	copyInputImage(imageMD);
	runInitialSegmentation();
	postProcess();
	setOutputImage();
	return &bw;
}

void Segmentation_GPU::displayFilteringResult(){
	float4* hostRGB = new float4[width * height];

	cudaMemcpy(hostRGB, clusterMap_C, width * height * sizeof(float4), cudaMemcpyDeviceToHost);
	cv::Mat_<cv::Vec3b> out(height, width);
	cv::Vec3b* ref;
	float4* devRef;
	for(int i = 0; i < height; i++){
		for(int j = 0; j < width; j++){
			ref = &out.at<cv::Vec3b>(cv::Point2d(j, i));
			devRef = &hostRGB[i * width + j];
			ref->val[0] = devRef->x;
			ref->val[1] = devRef->y;
			ref->val[2] = devRef->z;
		}
	}

	cv::cvtColor(out, out, CV_BGR2RGB);
	cv::imshow("filtered", out);
	cv::waitKey(1);

	delete[] hostRGB;
}


void Segmentation_GPU::copyInputImage(xn::ImageMetaData* imageMD){
	memcpy(original.data, imageMD->Data(), original.cols * original.rows * original.elemSize());
	cv::GaussianBlur(original, original, cv::Size(13, 13), 5.0);
	float4* temp = new float4[width * height];

	//int4.xyz = intensity
	//int4.w = cluster
	cv::Vec3b* c_l;
	for(int i = 0; i < height; i++)
		for(int j = 0; j < width; j++){
			c_l = &original.at<cv::Vec3b>(cv::Point2d(j, i));
			temp[i * width + j].x = c_l->val[0];
			temp[i * width + j].y = c_l->val[1];
			temp[i * width + j].z = c_l->val[2];
			temp[i * width + j].w = i * width + j;
		}

	cudaMemcpy(clusterMap_C, temp, width * height * sizeof(float4), cudaMemcpyHostToDevice);
	delete[] temp;
}

void Segmentation_GPU::setOutputImage(){
	float4* temp = new float4[width * height];
	cudaMemcpy(temp, clusterMap_C, width * height * sizeof(float4), cudaMemcpyDeviceToHost);
	for(int i = 0; i < height; i++)
		for(int j = 0; j < width; j++){
			bw.at<int>(cv::Point2d(j, i)) = temp[i * width + j].w;
			resultMap[i * width + j] = temp[i * width + j].w;
		}
		
	delete[] temp;	
}

int* Segmentation_GPU::getRawSegmentation(){
	return resultMap;

}









